{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a04d67a0-660f-41ea-a873-bab8c5f6197c",
   "metadata": {},
   "source": [
    "# Persist Messages\n",
    "\n",
    "When creating TableGPT agents, you have the option to persist their state, enabling interactions with the agent across multiple sessions while retaining memory of previous interactions. For more information on persistence, you can refer to the [Persistence](https://langchain-ai.github.io/langgraph/concepts/persistence/) documentation.\n",
    "\n",
    "The benefit of persistent messages is that you can interact with the TableGPT agent across multiple sessions, and the agent remembers previous interactions. This is useful for applications that require long-term tracking of context or complex conversations.\n",
    "\n",
    "TableGPT Agent leverages [langgraph-checkpoint](https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint) to implement persistent message storage. It supports using any type of `checkpointer` to store messages, such as: `Postgres`, `Redis`, `Memory`, etc. To integrate a checkpointer with a `TableGPT Agent`, you can follow the example below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7218312b-3482-44c1-afcd-48dafb48b83a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import date\n",
    "\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from tablegpt.agent import create_tablegpt_graph\n",
    "from pybox import LocalPyBoxManager\n",
    "\n",
    "llm = ChatOpenAI(openai_api_base=\"YOUR_VLLM_URL\", openai_api_key=\"whatever\", model_name=\"TableGPT2-7B\")\n",
    "pybox_manager = LocalPyBoxManager()\n",
    "checkpointer = MemorySaver()\n",
    "\n",
    "graph = create_tablegpt_graph(\n",
    "    llm=llm,\n",
    "    pybox_manager=pybox_manager,\n",
    "    checkpointer=checkpointer,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96b02c4f-d98b-42ca-a257-2db82c749f56",
   "metadata": {},
   "source": [
    "**Conducting a Conversation with `TableGPT`**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "841e7083-95aa-4e59-b66c-685acf2700f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"I understand that you're asking for an introduction to Jackie Chan. However, my primary role is to analyze datasets using Python. If you have a dataset related to Jackie Chan or any other topic, I'd be happy to help you analyze it. Could you please provide more details on what kind of data you have or what specific analysis you would like to perform?\", additional_kwargs={'parent_id': '1'}, response_metadata={}, id='cdf638ce-0e56-475b-a86b-0d8d7a0f6d05')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resp = await graph.ainvoke(\n",
    "    input={\n",
    "        \"messages\": [(\"human\", \"Please introduce Jackie Chan\")],\n",
    "        \"parent_id\": \"1\",\n",
    "        \"date\": date.today(),\n",
    "    },\n",
    "    config={\"configurable\": {\"thread_id\": \"1\"}},\n",
    ")\n",
    "resp[\"messages\"][-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c85dd8fd-36bc-42ab-b8d2-ed0f917093af",
   "metadata": {},
   "source": [
    "**Continuing the Conversation**\n",
    "\n",
    "To extend the conversation while maintaining context, you can provide new input along with the same `config` configuration:\n",
    "\n",
    "> Note: `config` is the configuration associated with this `checkpointer`. Through this configuration, the `checkpointer` can retrieve previous status information, so that in subsequent conversations, the model can better understand the user's intention and reply."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fe437db5-f6c9-40cc-a886-896e7e60ecad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"Certainly! Jackie Chan is a renowned actor, director, and martial artist, and he has starred in numerous films. Here are three popular movies in which he has participated:\\n\\n1. **Rush Hour (1998)** - In this action-comedy film, Jackie Chan plays the role of Inspector Lee, a Hong Kong detective who teams up with a Los Angeles detective, played by Chris Tucker, to solve a kidnapping case.\\n\\n2. **Drunken Master (1978)** - This is one of Jackie Chan's early films where he plays a young man who learns the art of drunken boxing to avenge his father's enemies.\\n\\n3. **The Karate Kid (2010)** - In this remake of the original 1984 film, Jackie Chan plays Mr. Han, a maintenance man who becomes the mentor to a young boy, Jaden Smith, teaching him martial arts and life lessons.\\n\\nIf you have any specific data or analysis related to these movies or Jackie Chan's filmography, feel free to provide more details, and I can help you with that!\", additional_kwargs={'parent_id': '1'}, response_metadata={}, id='4f62bec2-a4ec-43ad-97b2-cbade8c774b4')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resp = await graph.ainvoke(\n",
    "    input={\n",
    "        \"messages\": [(\"human\", \"Please name three movies he participated in.\")],\n",
    "        \"parent_id\": \"1\",\n",
    "        \"date\": date.today(),\n",
    "    },\n",
    "    config={\"configurable\": {\"thread_id\": \"1\"}},\n",
    ")\n",
    "resp[\"messages\"][-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac748856-f062-43a2-8381-6b80e7e2bf96",
   "metadata": {},
   "source": [
    "**Next, we demonstrates how to use `Postgres` as the backend for persisting checkpoint state using the [langgraph-checkpoint-postgres](https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-postgres) library.**\n",
    "\n",
    "## Installing Required Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6532184f-6871-4fb1-adb3-e61b1335cb9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U psycopg psycopg-pool psycopg_binary langgraph langgraph-checkpoint-postgres"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95b50676-b43b-429f-9c85-91b384832d60",
   "metadata": {},
   "source": [
    "## Use Async Connection\n",
    "\n",
    "**Note: `TableGPT Agent` is built based on [LangGraph](https://langchain-ai.github.io/langgraph/), and many of the `Node` components use `async/await` syntax, which does not yet support non-asynchronous operations.**\n",
    "\n",
    "Setting up an asynchronous connection to the database allows for non-blocking database operations. This means other parts of your application can continue running while waiting for database operations to complete. This is particularly beneficial in high-concurrency scenarios or when dealing with I/O-bound operations.\n",
    "\n",
    "The `DB_URI` is the database connection URI, specifying the protocol for connecting to a PostgreSQL database, including authentication and the host where the database is running."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dd0c3389-da09-41a5-ac70-d1ecd5bee211",
   "metadata": {},
   "outputs": [],
   "source": [
    "DB_URI = \"postgresql://postgres:postgres@127.0.0.1:5432/postgres?sslmode=disable\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "806b491f-f6e3-47d2-9043-23bf784635e1",
   "metadata": {},
   "source": [
    "### Creating a Checkpointer with AsyncPostgresSaver\n",
    "\n",
    "This creates a connection based on a connection string:\n",
    "- Advantages: Simplicity, encapsulates connection details\n",
    "- Best for: Quick setup or when connection details are provided as a string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3e78e56d-a108-4594-bd1d-d56189cfa9f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver\n",
    "from tablegpt.agent import create_tablegpt_graph\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
    "\n",
    "async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:\n",
    "    graph = create_tablegpt_graph(\n",
    "        llm=llm,\n",
    "        pybox_manager=pybox_manager,\n",
    "        checkpointer=checkpointer,\n",
    "    )\n",
    "    \n",
    "    res = await graph.ainvoke(\n",
    "        input={\n",
    "            \"messages\": [(\"human\", \"Who are you?\")],\n",
    "            \"parent_id\": \"2\",\n",
    "            \"date\": date.today()\n",
    "        },\n",
    "        config=config,\n",
    "    )\n",
    "    checkpoint_tuples = [c async for c in checkpointer.alist(config)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "18de503a-69f8-4191-9ed9-85bdd25d6213",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1efa6543-5d51-63c3-8001-9bd42cf9d6e6'}}, checkpoint={'v': 1, 'id': '1efa6543-5d51-63c3-8001-9bd42cf9d6e6', 'ts': '2024-11-19T08:56:48.239486+00:00', 'pending_sends': [], 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.6926057190269731'}, 'data_analyze_graph': {'branch:__start__:router:data_analyze_graph': '00000000000000000000000000000002.0.32407437506283565'}}, 'channel_versions': {'date': '00000000000000000000000000000003.0.1780977977687367', 'messages': '00000000000000000000000000000003.0.05509702188973753', '__start__': '00000000000000000000000000000002.3.0886787893869005e-05', 'parent_id': '00000000000000000000000000000003.0.43858547879187637', 'data_analyze_graph': '00000000000000000000000000000003.0.1082481333441786', 'branch:__start__:router:data_analyze_graph': '00000000000000000000000000000003.0.593567034515958'}, 'channel_values': {'date': datetime.date(2024, 11, 19), 'messages': [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930'), AIMessage(content=\"I am TableGPT2, an expert Python data analyst developed by Zhejiang University. My primary role is to assist you in analyzing datasets by writing Python code. I can help you with tasks such as data cleaning, transformation, visualization, and more. If you have a dataset or a specific analysis in mind, feel free to share it with me, and I'll do my best to help you!\", additional_kwargs={'parent_id': '2'}, response_metadata={}, id='560a88be-4fd0-4cc1-aa55-0747862fa222')], 'parent_id': '2', 'data_analyze_graph': 'data_analyze_graph'}}, metadata={'step': 1, 'source': 'loop', 'writes': {'data_analyze_graph': {'date': datetime.date(2024, 11, 19), 'messages': [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930'), AIMessage(content=\"I am TableGPT2, an expert Python data analyst developed by Zhejiang University. My primary role is to assist you in analyzing datasets by writing Python code. I can help you with tasks such as data cleaning, transformation, visualization, and more. If you have a dataset or a specific analysis in mind, feel free to share it with me, and I'll do my best to help you!\", additional_kwargs={'parent_id': '2'}, response_metadata={}, id='560a88be-4fd0-4cc1-aa55-0747862fa222')], 'parent_id': '2'}}, 'parents': {}, 'thread_id': '2'}, parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1efa6543-4a1a-6852-8000-b975c65fb2ff'}}, pending_writes=[]),\n",
       " CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'checkpoint_id': '1efa6543-5d27-6b37-8002-4e274906f4a5'}}, checkpoint={'v': 1, 'id': '1efa6543-5d27-6b37-8002-4e274906f4a5', 'ts': '2024-11-19T08:56:48.222457+00:00', 'pending_sends': [], 'versions_seen': {'agent': {'join:input_guard+retrieve_columns:agent': '00000000000000000000000000000003.0.8898909183470118'}, '__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.6163867462467301'}, 'input_guard': {'start:input_guard': '00000000000000000000000000000002.0.6848611387807798'}, 'retrieve_columns': {'start:retrieve_columns': '00000000000000000000000000000002.0.030416452982199194'}}, 'channel_versions': {'date': '00000000000000000000000000000002.0.2490273362085793', 'agent': '00000000000000000000000000000005.0.024232530645486583', 'messages': '00000000000000000000000000000005.0.14282746367420773', '__start__': '00000000000000000000000000000002.0.18620036399153372', 'parent_id': '00000000000000000000000000000002.0.5201095646733788', 'input_guard': '00000000000000000000000000000005.0.859473129275239', 'retrieve_columns': '00000000000000000000000000000005.0.7752176585300508', 'start:input_guard': '00000000000000000000000000000003.0.6183120254220215', 'start:retrieve_columns': '00000000000000000000000000000003.0.08187600687354024', 'join:input_guard+retrieve_columns:agent': '00000000000000000000000000000004.0.5107824581933167'}, 'channel_values': {'date': datetime.date(2024, 11, 19), 'agent': 'agent', 'messages': [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930'), AIMessage(content=\"I am TableGPT2, an expert Python data analyst developed by Zhejiang University. My primary role is to assist you in analyzing datasets by writing Python code. I can help you with tasks such as data cleaning, transformation, visualization, and more. If you have a dataset or a specific analysis in mind, feel free to share it with me, and I'll do my best to help you!\", additional_kwargs={'parent_id': '2'}, response_metadata={}, id='560a88be-4fd0-4cc1-aa55-0747862fa222')], 'parent_id': '2', 'join:input_guard+retrieve_columns:agent': set()}}, metadata={'step': 2, 'source': 'loop', 'writes': {'agent': {'messages': [AIMessage(content=\"I am TableGPT2, an expert Python data analyst developed by Zhejiang University. My primary role is to assist you in analyzing datasets by writing Python code. I can help you with tasks such as data cleaning, transformation, visualization, and more. If you have a dataset or a specific analysis in mind, feel free to share it with me, and I'll do my best to help you!\", additional_kwargs={'parent_id': '2'}, response_metadata={}, id='560a88be-4fd0-4cc1-aa55-0747862fa222')]}}, 'parents': {'': '1efa6543-4a1a-6852-8000-b975c65fb2ff'}, 'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'langgraph_node': 'data_analyze_graph', 'langgraph_path': ['__pregel_pull', 'data_analyze_graph'], 'langgraph_step': 1, 'langgraph_triggers': ['branch:__start__:router:data_analyze_graph'], 'langgraph_checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd'}, parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'checkpoint_id': '1efa6543-4a53-6aa4-8001-7474db32528e'}}, pending_writes=[]),\n",
       " CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'checkpoint_id': '1efa6543-4a53-6aa4-8001-7474db32528e'}}, checkpoint={'v': 1, 'id': '1efa6543-4a53-6aa4-8001-7474db32528e', 'ts': '2024-11-19T08:56:46.248182+00:00', 'pending_sends': [], 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.6163867462467301'}, 'input_guard': {'start:input_guard': '00000000000000000000000000000002.0.6848611387807798'}, 'retrieve_columns': {'start:retrieve_columns': '00000000000000000000000000000002.0.030416452982199194'}}, 'channel_versions': {'date': '00000000000000000000000000000002.0.2490273362085793', 'messages': '00000000000000000000000000000003.0.8478737633204881', '__start__': '00000000000000000000000000000002.0.18620036399153372', 'parent_id': '00000000000000000000000000000002.0.5201095646733788', 'input_guard': '00000000000000000000000000000003.0.06039244147872136', 'retrieve_columns': '00000000000000000000000000000003.0.8552403538042089', 'start:input_guard': '00000000000000000000000000000003.0.6183120254220215', 'start:retrieve_columns': '00000000000000000000000000000003.0.08187600687354024', 'join:input_guard+retrieve_columns:agent': '00000000000000000000000000000003.0.8898909183470118'}, 'channel_values': {'date': datetime.date(2024, 11, 19), 'messages': [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930')], 'parent_id': '2', 'input_guard': 'input_guard', 'retrieve_columns': 'retrieve_columns', 'join:input_guard+retrieve_columns:agent': {'input_guard', 'retrieve_columns'}}}, metadata={'step': 1, 'source': 'loop', 'writes': {'input_guard': {'messages': []}, 'retrieve_columns': {'messages': []}}, 'parents': {'': '1efa6543-4a1a-6852-8000-b975c65fb2ff'}, 'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'langgraph_node': 'data_analyze_graph', 'langgraph_path': ['__pregel_pull', 'data_analyze_graph'], 'langgraph_step': 1, 'langgraph_triggers': ['branch:__start__:router:data_analyze_graph'], 'langgraph_checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd'}, parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'checkpoint_id': '1efa6543-4a4c-6639-8000-a5baf4d38bd2'}}, pending_writes=[('36d7d700-5338-feda-05f0-57d74fddbc0b', 'agent', 'agent'), ('36d7d700-5338-feda-05f0-57d74fddbc0b', 'messages', [AIMessage(content=\"I am TableGPT2, an expert Python data analyst developed by Zhejiang University. My primary role is to assist you in analyzing datasets by writing Python code. I can help you with tasks such as data cleaning, transformation, visualization, and more. If you have a dataset or a specific analysis in mind, feel free to share it with me, and I'll do my best to help you!\", additional_kwargs={'parent_id': '2'}, response_metadata={}, id='560a88be-4fd0-4cc1-aa55-0747862fa222')])]),\n",
       " CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'checkpoint_id': '1efa6543-4a4c-6639-8000-a5baf4d38bd2'}}, checkpoint={'v': 1, 'id': '1efa6543-4a4c-6639-8000-a5baf4d38bd2', 'ts': '2024-11-19T08:56:46.245208+00:00', 'pending_sends': [], 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.6163867462467301'}}, 'channel_versions': {'date': '00000000000000000000000000000002.0.2490273362085793', 'messages': '00000000000000000000000000000002.0.19507603965774079', '__start__': '00000000000000000000000000000002.0.18620036399153372', 'parent_id': '00000000000000000000000000000002.0.5201095646733788', 'start:input_guard': '00000000000000000000000000000002.0.6848611387807798', 'start:retrieve_columns': '00000000000000000000000000000002.0.030416452982199194'}, 'channel_values': {'date': datetime.date(2024, 11, 19), 'messages': [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930')], 'parent_id': '2', 'start:input_guard': '__start__', 'start:retrieve_columns': '__start__'}}, metadata={'step': 0, 'source': 'loop', 'writes': None, 'parents': {'': '1efa6543-4a1a-6852-8000-b975c65fb2ff'}, 'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'langgraph_node': 'data_analyze_graph', 'langgraph_path': ['__pregel_pull', 'data_analyze_graph'], 'langgraph_step': 1, 'langgraph_triggers': ['branch:__start__:router:data_analyze_graph'], 'langgraph_checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd'}, parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'checkpoint_id': '1efa6543-4a49-6ec0-bfff-7a6bdf830d6b'}}, pending_writes=[('637b8d5c-1e9a-d15c-7560-eba440c88860', 'input_guard', 'input_guard'), ('637b8d5c-1e9a-d15c-7560-eba440c88860', 'messages', []), ('637b8d5c-1e9a-d15c-7560-eba440c88860', 'join:input_guard+retrieve_columns:agent', 'input_guard'), ('fcc182eb-567e-cef1-c5be-16b527e21434', 'retrieve_columns', 'retrieve_columns'), ('fcc182eb-567e-cef1-c5be-16b527e21434', 'messages', []), ('fcc182eb-567e-cef1-c5be-16b527e21434', 'join:input_guard+retrieve_columns:agent', 'retrieve_columns')]),\n",
       " CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'checkpoint_id': '1efa6543-4a49-6ec0-bfff-7a6bdf830d6b'}}, checkpoint={'v': 1, 'id': '1efa6543-4a49-6ec0-bfff-7a6bdf830d6b', 'ts': '2024-11-19T08:56:46.244206+00:00', 'pending_sends': [], 'versions_seen': {'__input__': {}}, 'channel_versions': {'__start__': '00000000000000000000000000000001.0.6163867462467301'}, 'channel_values': {'__start__': {'messages': [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930')], 'parent_id': '2', 'date': datetime.date(2024, 11, 19)}}}, metadata={'step': -1, 'source': 'input', 'writes': {'__start__': {'date': datetime.date(2024, 11, 19), 'messages': [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930')], 'parent_id': '2'}}, 'parents': {'': '1efa6543-4a1a-6852-8000-b975c65fb2ff'}, 'thread_id': '2', 'checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'langgraph_node': 'data_analyze_graph', 'langgraph_path': ['__pregel_pull', 'data_analyze_graph'], 'langgraph_step': 1, 'langgraph_triggers': ['branch:__start__:router:data_analyze_graph'], 'langgraph_checkpoint_ns': 'data_analyze_graph:3d6fb60f-4da5-a1de-bcf2-fa5632547abd'}, parent_config=None, pending_writes=[('39da96de-984e-3d02-e2ef-9bd5146d7336', 'messages', [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930')]), ('39da96de-984e-3d02-e2ef-9bd5146d7336', 'date', datetime.date(2024, 11, 19)), ('39da96de-984e-3d02-e2ef-9bd5146d7336', 'parent_id', '2'), ('39da96de-984e-3d02-e2ef-9bd5146d7336', 'start:input_guard', '__start__'), ('39da96de-984e-3d02-e2ef-9bd5146d7336', 'start:retrieve_columns', '__start__')]),\n",
       " CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1efa6543-4a1a-6852-8000-b975c65fb2ff'}}, checkpoint={'v': 1, 'id': '1efa6543-4a1a-6852-8000-b975c65fb2ff', 'ts': '2024-11-19T08:56:46.224784+00:00', 'pending_sends': [], 'versions_seen': {'__input__': {}, '__start__': {'__start__': '00000000000000000000000000000001.0.6926057190269731'}}, 'channel_versions': {'date': '00000000000000000000000000000002.0.602279509708772', 'messages': '00000000000000000000000000000002.0.47212683047327253', '__start__': '00000000000000000000000000000002.3.0886787893869005e-05', 'parent_id': '00000000000000000000000000000002.0.43326965052986344', 'branch:__start__:router:data_analyze_graph': '00000000000000000000000000000002.0.32407437506283565'}, 'channel_values': {'date': datetime.date(2024, 11, 19), 'messages': [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930')], 'parent_id': '2', 'branch:__start__:router:data_analyze_graph': '__start__'}}, metadata={'step': 0, 'source': 'loop', 'writes': None, 'parents': {}, 'thread_id': '2'}, parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1efa6543-4a12-6bbb-bfff-ac4aee846bfe'}}, pending_writes=[('3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'data_analyze_graph', 'data_analyze_graph'), ('3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'messages', [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930'), AIMessage(content=\"I am TableGPT2, an expert Python data analyst developed by Zhejiang University. My primary role is to assist you in analyzing datasets by writing Python code. I can help you with tasks such as data cleaning, transformation, visualization, and more. If you have a dataset or a specific analysis in mind, feel free to share it with me, and I'll do my best to help you!\", additional_kwargs={'parent_id': '2'}, response_metadata={}, id='560a88be-4fd0-4cc1-aa55-0747862fa222')]), ('3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'parent_id', '2'), ('3d6fb60f-4da5-a1de-bcf2-fa5632547abd', 'date', datetime.date(2024, 11, 19))]),\n",
       " CheckpointTuple(config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1efa6543-4a12-6bbb-bfff-ac4aee846bfe'}}, checkpoint={'v': 1, 'id': '1efa6543-4a12-6bbb-bfff-ac4aee846bfe', 'ts': '2024-11-19T08:56:46.221598+00:00', 'pending_sends': [], 'versions_seen': {'__input__': {}}, 'channel_versions': {'__start__': '00000000000000000000000000000001.0.6926057190269731'}, 'channel_values': {'__start__': {'messages': [['human', 'Who are you?']], 'parent_id': '2', 'date': datetime.date(2024, 11, 19)}}}, metadata={'step': -1, 'source': 'input', 'writes': {'__start__': {'date': datetime.date(2024, 11, 19), 'messages': [['human', 'Who are you?']], 'parent_id': '2'}}, 'parents': {}, 'thread_id': '2'}, parent_config=None, pending_writes=[('8e27b3ac-0a9d-697e-0667-6d3ccc170a50', 'messages', [['human', 'Who are you?']]), ('8e27b3ac-0a9d-697e-0667-6d3ccc170a50', 'parent_id', '2'), ('8e27b3ac-0a9d-697e-0667-6d3ccc170a50', 'date', datetime.date(2024, 11, 19)), ('8e27b3ac-0a9d-697e-0667-6d3ccc170a50', 'branch:__start__:router:data_analyze_graph', '__start__')])]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "checkpoint_tuples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c3d22bf-ebdb-48fe-a452-8697d786404e",
   "metadata": {},
   "source": [
    "## Get Persisted Messages with Config\n",
    "\n",
    "We can use the same config parameters to retrieve persisted messages through the `checkpointer`. You can follow the example below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c1b166b5-cbfb-4cbe-b193-605f64315a38",
   "metadata": {},
   "outputs": [],
   "source": [
    "async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:\n",
    "    graph = create_tablegpt_graph(\n",
    "        llm=llm,\n",
    "        pybox_manager=pybox_manager,\n",
    "        checkpointer=checkpointer,\n",
    "    )\n",
    "    \n",
    "    graph_state = await graph.aget_state(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "09963483-7f79-4050-a27d-3862b0b940e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StateSnapshot(values={'messages': [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930'), AIMessage(content=\"I am TableGPT2, an expert Python data analyst developed by Zhejiang University. My primary role is to assist you in analyzing datasets by writing Python code. I can help you with tasks such as data cleaning, transformation, visualization, and more. If you have a dataset or a specific analysis in mind, feel free to share it with me, and I'll do my best to help you!\", additional_kwargs={'parent_id': '2'}, response_metadata={}, id='560a88be-4fd0-4cc1-aa55-0747862fa222')], 'parent_id': '2', 'date': datetime.date(2024, 11, 19)}, next=(), config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1efa6543-5d51-63c3-8001-9bd42cf9d6e6'}}, metadata={'step': 1, 'source': 'loop', 'writes': {'data_analyze_graph': {'date': datetime.date(2024, 11, 19), 'messages': [HumanMessage(content='Who are you?', additional_kwargs={}, response_metadata={}, id='32b67f59-d13e-43eb-9239-ec711811e930'), AIMessage(content=\"I am TableGPT2, an expert Python data analyst developed by Zhejiang University. My primary role is to assist you in analyzing datasets by writing Python code. I can help you with tasks such as data cleaning, transformation, visualization, and more. If you have a dataset or a specific analysis in mind, feel free to share it with me, and I'll do my best to help you!\", additional_kwargs={'parent_id': '2'}, response_metadata={}, id='560a88be-4fd0-4cc1-aa55-0747862fa222')], 'parent_id': '2'}}, 'parents': {}, 'thread_id': '2'}, created_at='2024-11-19T08:56:48.239486+00:00', parent_config={'configurable': {'thread_id': '2', 'checkpoint_ns': '', 'checkpoint_id': '1efa6543-4a1a-6852-8000-b975c65fb2ff'}}, tasks=())"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph_state"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
